{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hBgIy9eYYx35"
      },
      "source": [
        "## **Installation**\n",
        "Requirements:\n",
        "- python >= 3.9\n",
        "- a device 2 GPUs\n",
        "\n",
        "We highly recommend CUDA when using torchRec. If using CUDA:\n",
        "- cuda >= 12.0\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Akmt4viCo9dz",
        "outputId": "e008352d-85b6-4713-827f-7a3eeb9ad09b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Looking in indexes: https://download.pytorch.org/whl/cu121\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (2.5.1+cu121)\n",
            "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.10/dist-packages (from torch) (1.13.1)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch) (3.3)\n",
            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch) (12.1.105)\n",
            "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.10/dist-packages (from torch) (4.13.2)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch) (2024.6.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch) (3.1.6)\n",
            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /usr/local/lib/python3.10/dist-packages (from torch) (12.1.3.1)\n",
            "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /usr/local/lib/python3.10/dist-packages (from torch) (9.1.0.70)\n",
            "Requirement already satisfied: triton==3.1.0 in /usr/local/lib/python3.10/dist-packages (from torch) (3.1.0)\n",
            "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.10/dist-packages (from torch) (2.21.5)\n",
            "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /usr/local/lib/python3.10/dist-packages (from torch) (11.0.2.54)\n",
            "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /usr/local/lib/python3.10/dist-packages (from torch) (12.1.0.106)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch) (3.13.1)\n",
            "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /usr/local/lib/python3.10/dist-packages (from torch) (10.3.2.106)\n",
            "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /usr/local/lib/python3.10/dist-packages (from torch) (11.4.5.107)\n",
            "Requirement already satisfied: nvidia-nvjitlink-cu12 in /usr/local/lib/python3.10/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->torch) (12.1.105)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy==1.13.1->torch) (1.3.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch) (3.0.2)\n",
            "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
            "\u001b[0mLooking in indexes: https://download.pytorch.org/whl/cu121\n",
            "Requirement already satisfied: fbgemm_gpu in /usr/local/lib/python3.10/dist-packages (1.0.0+cu121)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from fbgemm_gpu) (2.2.5)\n",
            "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
            "\u001b[0mRequirement already satisfied: torchmetrics in /usr/local/lib/python3.10/dist-packages (1.0.3)\n",
            "Requirement already satisfied: lightning-utilities>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from torchmetrics) (0.14.3)\n",
            "Requirement already satisfied: torch>=1.8.1 in /usr/local/lib/python3.10/dist-packages (from torchmetrics) (2.5.1+cu121)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from torchmetrics) (25.0)\n",
            "Requirement already satisfied: numpy>1.20.0 in /usr/local/lib/python3.10/dist-packages (from torchmetrics) (2.2.5)\n",
            "Requirement already satisfied: typing_extensions in /usr/local/lib/python3.10/dist-packages (from lightning-utilities>=0.7.0->torchmetrics) (4.13.2)\n",
            "Requirement already satisfied: setuptools in /usr/lib/python3/dist-packages (from lightning-utilities>=0.7.0->torchmetrics) (59.6.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (3.13.1)\n",
            "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (12.1.105)\n",
            "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (2.21.5)\n",
            "Requirement already satisfied: triton==3.1.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (3.1.0)\n",
            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (12.1.105)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (2024.6.1)\n",
            "Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (9.1.0.70)\n",
            "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (12.1.3.1)\n",
            "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (1.13.1)\n",
            "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (10.3.2.106)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (3.1.6)\n",
            "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (11.0.2.54)\n",
            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (11.4.5.107)\n",
            "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (12.1.0.106)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (3.3)\n",
            "Requirement already satisfied: nvidia-nvjitlink-cu12 in /usr/local/lib/python3.10/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->torch>=1.8.1->torchmetrics) (12.1.105)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy==1.13.1->torch>=1.8.1->torchmetrics) (1.3.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.8.1->torchmetrics) (3.0.2)\n",
            "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
            "\u001b[0mLooking in indexes: https://download.pytorch.org/whl/cu121\n",
            "Requirement already satisfied: torchrec in /usr/local/lib/python3.10/dist-packages (1.0.0+cu121)\n",
            "Requirement already satisfied: fbgemm-gpu in /usr/local/lib/python3.10/dist-packages (from torchrec) (1.0.0+cu121)\n",
            "Requirement already satisfied: torchmetrics==1.0.3 in /usr/local/lib/python3.10/dist-packages (from torchrec) (1.0.3)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from torchrec) (4.66.5)\n",
            "Requirement already satisfied: pyre-extensions in /usr/local/lib/python3.10/dist-packages (from torchrec) (0.0.31)\n",
            "Requirement already satisfied: iopath in /usr/local/lib/python3.10/dist-packages (from torchrec) (0.1.9)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from torchmetrics==1.0.3->torchrec) (25.0)\n",
            "Requirement already satisfied: torch>=1.8.1 in /usr/local/lib/python3.10/dist-packages (from torchmetrics==1.0.3->torchrec) (2.5.1+cu121)\n",
            "Requirement already satisfied: numpy>1.20.0 in /usr/local/lib/python3.10/dist-packages (from torchmetrics==1.0.3->torchrec) (2.2.5)\n",
            "Requirement already satisfied: lightning-utilities>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from torchmetrics==1.0.3->torchrec) (0.14.3)\n",
            "Requirement already satisfied: portalocker in /usr/local/lib/python3.10/dist-packages (from iopath->torchrec) (2.10.1)\n",
            "Requirement already satisfied: typing-inspect in /usr/local/lib/python3.10/dist-packages (from pyre-extensions->torchrec) (0.9.0)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from pyre-extensions->torchrec) (4.13.2)\n",
            "Requirement already satisfied: setuptools in /usr/lib/python3/dist-packages (from lightning-utilities>=0.7.0->torchmetrics==1.0.3->torchrec) (59.6.0)\n",
            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (11.4.5.107)\n",
            "Requirement already satisfied: triton==3.1.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (3.1.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (3.13.1)\n",
            "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (1.13.1)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (3.3)\n",
            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (9.1.0.70)\n",
            "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (12.1.0.106)\n",
            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (12.1.3.1)\n",
            "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (11.0.2.54)\n",
            "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (10.3.2.106)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (2024.6.1)\n",
            "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (2.21.5)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics==1.0.3->torchrec) (3.1.6)\n",
            "Requirement already satisfied: nvidia-nvjitlink-cu12 in /usr/local/lib/python3.10/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->torch>=1.8.1->torchmetrics==1.0.3->torchrec) (12.1.105)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy==1.13.1->torch>=1.8.1->torchmetrics==1.0.3->torchrec) (1.3.0)\n",
            "Requirement already satisfied: mypy-extensions>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from typing-inspect->pyre-extensions->torchrec) (1.0.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.8.1->torchmetrics==1.0.3->torchrec) (3.0.2)\n",
            "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
            "\u001b[0m"
          ]
        }
      ],
      "source": [
        "!pip3 install --pre torch --index-url https://download.pytorch.org/whl/cu121 -U\n",
        "!pip3 install fbgemm_gpu --index-url https://download.pytorch.org/whl/cu121\n",
        "!pip3 install torchmetrics\n",
        "!pip3 install torchrec --index-url https://download.pytorch.org/whl/cu121"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pkUJVDAA2dER",
        "outputId": "2ffbbf14-b6b4-4687-9b0d-4bba85d8c316"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: multiprocess in /usr/local/lib/python3.10/dist-packages (0.70.18)\n",
            "Requirement already satisfied: dill>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from multiprocess) (0.4.0)\n",
            "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
            "\u001b[0m"
          ]
        }
      ],
      "source": [
        "!pip3 install multiprocess"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HWBOrwVSnrNE"
      },
      "source": [
        "## **Overview**\n",
        "This tutorial will mainly cover the sharding schemes of embedding tables via `EmbeddingPlanner` and `DistributedModelParallel` API and explore the benefits of different sharding schemes for the embedding tables by explicitly configuring them."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "udsN6PlUo1zF"
      },
      "source": [
        "### Distributed Setup\n",
        "Due to the notebook enviroment, we cannot run [`SPMD`](https://en.wikipedia.org/wiki/SPMD) program here but we can do multiprocessing inside the notebook to mimic the setup. Users should be responsible for setting up their own [`SPMD`](https://en.wikipedia.org/wiki/SPMD) launcher when using Torchrec.\n",
        "We setup our environment so that torch distributed based communication backend can work."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "4-v17rxkopQw"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import copy\n",
        "import torch\n",
        "import torchrec\n",
        "import multiprocess\n",
        "from torchrec.distributed.types import ShardingEnv\n",
        "\n",
        "os.environ[\"MASTER_ADDR\"] = \"localhost\"\n",
        "os.environ[\"MASTER_PORT\"] = \"10000\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "td6bvF_KRbzx"
      },
      "source": [
        "Below are codes setup for one process at rank `0` and with WORLD_SIZE (number of processes) as `1`.\n",
        "\n",
        "In a distributed setup, we will repeat the following steps on each process to set up the distributed environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "suuMmJs2RbjX"
      },
      "outputs": [],
      "source": [
        "import torch.distributed as dist\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "rank = 0\n",
        "world_size = 1\n",
        "backend = 'nccl' if 'cuda' in device else 'gloo'\n",
        "os.environ[\"RANK\"] = f\"{rank}\"\n",
        "os.environ[\"WORLD_SIZE\"] = f\"{world_size}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZdSUWBRxoP8R"
      },
      "source": [
        "### Constructing our embedding model\n",
        "Here we use TorchRec offering of [`EmbeddingBagCollection`](https://github.com/facebookresearch/torchrec/blob/main/torchrec/modules/embedding_modules.py#L59) to construct our embedding bag model with embedding tables.\n",
        "\n",
        "Here, we create an EmbeddingBagCollection (EBC) with four embedding bags. We have two types of tables: large tables and small tables differentiated by their row size difference: 4096 vs 1024. Each table is still represented by 64 dimension embedding.\n",
        "\n",
        "We configure the `ParameterConstraints` data structure for the tables, which provides hints for the model parallel API to help decide the sharding and placement strategy for the tables.\n",
        "In TorchRec, we support\n",
        "* `table-wise`: place the entire table on one device;\n",
        "* `row-wise`: shard the table evenly by row dimension and place one shard on each device of the communication world;\n",
        "* `column-wise`: shard the table evenly by embedding dimension, and place one shard on each device of the communication world;\n",
        "* `table-row-wise`: special sharding optimized for intra-host communication for available fast intra-machine device interconnect, e.g. NVLink;\n",
        "* `data_parallel`: replicate the tables for every device;\n",
        "\n",
        "Note how we initially allocate the EBC on device \"meta\". This will tell EBC to not allocate memory yet."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "e7UQBuG09hbj"
      },
      "outputs": [],
      "source": [
        "from torchrec.distributed.planner.types import ParameterConstraints\n",
        "from torchrec.distributed.embedding_types import EmbeddingComputeKernel\n",
        "from torchrec.distributed.types import ShardingType, ShardingPlan\n",
        "from typing import Dict\n",
        "\n",
        "large_table_cnt = 2\n",
        "small_table_cnt = 2\n",
        "large_tables=[\n",
        "  torchrec.EmbeddingBagConfig(\n",
        "    name=\"large_table_\" + str(i),\n",
        "    embedding_dim=64,\n",
        "    num_embeddings=4096,\n",
        "    feature_names=[\"large_table_feature_\" + str(i)],\n",
        "    pooling=torchrec.PoolingType.SUM,\n",
        "  ) for i in range(large_table_cnt)\n",
        "]\n",
        "small_tables=[\n",
        "  torchrec.EmbeddingBagConfig(\n",
        "    name=\"small_table_\" + str(i),\n",
        "    embedding_dim=64,\n",
        "    num_embeddings=1024,\n",
        "    feature_names=[\"small_table_feature_\" + str(i)],\n",
        "    pooling=torchrec.PoolingType.SUM,\n",
        "  ) for i in range(small_table_cnt)\n",
        "]\n",
        "\n",
        "def gen_constraints(sharding_type: ShardingType = ShardingType.TABLE_WISE) -> Dict[str, ParameterConstraints]:\n",
        "  large_table_constraints = {\n",
        "    \"large_table_\" + str(i): ParameterConstraints(\n",
        "      sharding_types=[sharding_type.value],\n",
        "    ) for i in range(large_table_cnt)\n",
        "  }\n",
        "  small_table_constraints = {\n",
        "    \"small_table_\" + str(i): ParameterConstraints(\n",
        "      sharding_types=[sharding_type.value],\n",
        "    ) for i in range(small_table_cnt)\n",
        "  }\n",
        "  constraints = {**large_table_constraints, **small_table_constraints}\n",
        "  return constraints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "U0ktiI29_FVW"
      },
      "outputs": [],
      "source": [
        "ebc = torchrec.EmbeddingBagCollection(\n",
        "    device=torch.device(device),\n",
        "    tables=large_tables + small_tables\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UMoH-xpV2S5w",
        "outputId": "0aa9ca26-2940-479f-c3ee-fef10868258a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "EmbeddingBagCollection(\n",
            "  (embedding_bags): ModuleDict(\n",
            "    (large_table_0): EmbeddingBag(4096, 64, mode='sum')\n",
            "    (large_table_1): EmbeddingBag(4096, 64, mode='sum')\n",
            "    (small_table_0): EmbeddingBag(1024, 64, mode='sum')\n",
            "    (small_table_1): EmbeddingBag(1024, 64, mode='sum')\n",
            "  )\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "print(ebc)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "711VBygVHGJ6"
      },
      "source": [
        "For `table-row-wise`, unfortuately we cannot simulate it due to its nature of operating under multi-host setup. We will present a python [`SPMD`](https://en.wikipedia.org/wiki/SPMD) example in the future to train models with `table-row-wise`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1G8aUfmeMA7m"
      },
      "source": [
        "\n",
        "With data parallel, we will repeat the tables for all devices.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7m0_ssVLFQEH"
      },
      "source": [
        "### DistributedModelParallel in multiprocessing\n",
        "If you have access for **2 GPUs**, we can work on multi-GPU multi-process sharding. Though due to the issue in \"Spawn\"-started multiprocess, the print may not have outputs on certain devices. But you can check if the assertion is passed.\n",
        "\n",
        "we have a single process execution function for mimicking one rank's work during [`SPMD`](https://en.wikipedia.org/wiki/SPMD) execution.\n",
        "\n",
        "This code will shard the model collectively with other processes and allocate memories accordingly. It first sets up process groups and do embedding table placement using planner and generate sharded model using `DistributedModelParallel`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "PztCaGmLA85u"
      },
      "outputs": [],
      "source": [
        "def single_rank_execution(\n",
        "    rank: int,\n",
        "    world_size: int,\n",
        "    constraints: Dict[str, ParameterConstraints],\n",
        "    module: torch.nn.Module,\n",
        "    backend: str\n",
        ") -> None:\n",
        "\n",
        "    import os\n",
        "    import torch\n",
        "    import torch.distributed as dist\n",
        "    from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder\n",
        "    from torchrec.distributed.model_parallel import DistributedModelParallel\n",
        "    from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology\n",
        "    from torchrec.distributed.types import ModuleSharder, ShardingEnv, ShardingPlan\n",
        "    from typing import cast\n",
        "\n",
        "    def init_distributed_single_host(\n",
        "        rank: int,\n",
        "        world_size: int,\n",
        "        backend: str,\n",
        "        # pyre-fixme[11]: Annotation `ProcessGroup` is not defined as a type.\n",
        "    ) -> dist.ProcessGroup:\n",
        "        os.environ[\"RANK\"] = f\"{rank}\"\n",
        "        os.environ[\"WORLD_SIZE\"] = f\"{world_size}\"\n",
        "        dist.init_process_group(rank=rank, world_size=world_size, backend=backend)\n",
        "        return dist.group.WORLD\n",
        "\n",
        "    if backend == \"nccl\":\n",
        "        device = torch.device(f\"cuda:{rank}\")\n",
        "        torch.cuda.set_device(device)\n",
        "    else:\n",
        "        device = torch.device(\"cpu\")\n",
        "    topology = Topology(world_size=world_size, compute_device=device.type)\n",
        "    pg = init_distributed_single_host(rank, world_size, backend)\n",
        "    # pg = dist.group.WORLD\n",
        "    planner = EmbeddingShardingPlanner(\n",
        "        topology=topology,\n",
        "        constraints=constraints,\n",
        "    )\n",
        "    sharders = [cast(ModuleSharder[torch.nn.Module], EmbeddingBagCollectionSharder())]\n",
        "    plan: ShardingPlan = planner.collective_plan(module, sharders, pg)\n",
        "\n",
        "    sharded_model = DistributedModelParallel(\n",
        "        module,\n",
        "        env=ShardingEnv.from_process_group(pg),\n",
        "        plan=plan,\n",
        "        sharders=sharders,\n",
        "        device=device,\n",
        "    )\n",
        "    print(f\"rank:{rank},sharding plan: {plan}\")\n",
        "\n",
        "    return sharded_model\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3YvDnV_wz_An"
      },
      "source": [
        "### Multiprocessing Execution\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "arW0Jf6qEl-h"
      },
      "outputs": [],
      "source": [
        "def spmd_sharing_simulation(\n",
        "    sharding_type: ShardingType = ShardingType.TABLE_WISE,\n",
        "    world_size = 2, # Change this world size according to the number of GPUs available on your end.\n",
        "):\n",
        "  ctx = multiprocess.get_context(\"spawn\")\n",
        "  processes = []\n",
        "\n",
        "  for rank in range(world_size):\n",
        "      p = ctx.Process(\n",
        "          target=single_rank_execution,\n",
        "          args=(\n",
        "              rank,\n",
        "              world_size,\n",
        "              gen_constraints(sharding_type),\n",
        "              ebc,\n",
        "              backend,\n",
        "          ),\n",
        "      )\n",
        "      print(f\"start for rank: {rank}\")\n",
        "      p.start()\n",
        "      processes.append(p)\n",
        "\n",
        "  for p in processes:\n",
        "      p.join()\n",
        "      print(f\"exit code: {p.exitcode}\")\n",
        "      assert 0 == p.exitcode"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iUFvR2AeQQYe"
      },
      "source": [
        "Now we can start the multiprocess sharding. There will not be any output in the terminal due to the Spawn method used for generating child processes, but you should see all the assertions will pass and processes exit with code 0."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ux63HCxRN5Xr"
      },
      "source": [
        "### Table Wise Sharding\n",
        "Now let's execute the code in two processes for 2 GPUs. We can see in the plan print that how our tables are sharded across GPUs. Each node will have one large table and one small which shows our planner tries for load balance for the embedding tables. Table-wise is the de-factor go-to sharding schemes for many small-medium size tables for load balancing over the devices."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 263
        },
        "id": "gZfZatL5QJqc",
        "outputId": "ed6c86b8-4e8d-418d-e5be-089a12e98f89"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "start for rank: 0\n",
            "start for rank: 1\n",
            "rank:0,sharding plan: module: \n",
            "\n",
            "    param     | sharding type | compute kernel | ranks\n",
            "------------- | ------------- | -------------- | -----\n",
            "large_table_0 | table_wise    | fused          | [0]  \n",
            "large_table_1 | table_wise    | fused          | [1]  \n",
            "small_table_0 | table_wise    | fused          | [0]  \n",
            "small_table_1 | table_wise    | fused          | [1]  \n",
            "\n",
            "    param     | shard offsets | shard sizes |   placement  \n",
            "------------- | ------------- | ----------- | -------------\n",
            "large_table_0 | [0, 0]        | [4096, 64]  | rank:0/cuda:0\n",
            "large_table_1 | [0, 0]        | [4096, 64]  | rank:1/cuda:1\n",
            "small_table_0 | [0, 0]        | [1024, 64]  | rank:0/cuda:0\n",
            "small_table_1 | [0, 0]        | [1024, 64]  | rank:1/cuda:1\n",
            "rank:1,sharding plan: module: \n",
            "\n",
            "    param     | sharding type | compute kernel | ranks\n",
            "------------- | ------------- | -------------- | -----\n",
            "large_table_0 | table_wise    | fused          | [0]  \n",
            "large_table_1 | table_wise    | fused          | [1]  \n",
            "small_table_0 | table_wise    | fused          | [0]  \n",
            "small_table_1 | table_wise    | fused          | [1]  \n",
            "\n",
            "    param     | shard offsets | shard sizes |   placement  \n",
            "------------- | ------------- | ----------- | -------------\n",
            "large_table_0 | [0, 0]        | [4096, 64]  | rank:0/cuda:0\n",
            "large_table_1 | [0, 0]        | [4096, 64]  | rank:1/cuda:1\n",
            "small_table_0 | [0, 0]        | [1024, 64]  | rank:0/cuda:0\n",
            "small_table_1 | [0, 0]        | [1024, 64]  | rank:1/cuda:1\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "[rank0]:[W523 02:45:21.010274288 ProcessGroupNCCL.cpp:1250] Warning: WARNING: process group has NOT been destroyed before we destruct ProcessGroupNCCL. On normal program exit, the application should call destroy_process_group to ensure that any pending NCCL operations have finished in this process. In rare cases this process can exit before this point and block the progress of another member of the process group. This constraint has always been present,  but this warning has only been added since PyTorch 2.4 (function operator())\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "exit code: 0\n",
            "exit code: 0\n"
          ]
        }
      ],
      "source": [
        "spmd_sharing_simulation(ShardingType.TABLE_WISE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kfiXoz-2NBj4"
      },
      "source": [
        "### Explore other sharding modes\n",
        "We have initially explored what table-wise sharding would look like and how it balances the tables placement. Now we explore sharding modes with finer focus on load balance: row-wise.\n",
        "\n",
        "Row-wise is specifically addressing large tables which a single device cannot hold due to the memory size increase from large embedding row numbers. It can address the placement of the super large tables in your models. Users can see that in the `shard_sizes` section in the printed plan log, the tables are halved by row dimension to be distributed onto two GPUs.\n",
        "\n",
        "If you are on a CPU, the row-wise sharding will not be allowed.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "853VYV_0P3us"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "start for rank: 0\n",
            "start for rank: 1\n",
            "rank:0,sharding plan: module: \n",
            "\n",
            "    param     | sharding type | compute kernel | ranks \n",
            "------------- | ------------- | -------------- | ------\n",
            "large_table_0 | row_wise      | fused          | [0, 1]\n",
            "large_table_1 | row_wise      | fused          | [0, 1]\n",
            "small_table_0 | row_wise      | fused          | [0, 1]\n",
            "small_table_1 | row_wise      | fused          | [0, 1]\n",
            "\n",
            "    param     | shard offsets | shard sizes |   placement  \n",
            "------------- | ------------- | ----------- | -------------\n",
            "large_table_0 | [0, 0]        | [2048, 64]  | rank:0/cuda:0\n",
            "large_table_0 | [2048, 0]     | [2048, 64]  | rank:1/cuda:1\n",
            "large_table_1 | [0, 0]        | [2048, 64]  | rank:0/cuda:0\n",
            "large_table_1 | [2048, 0]     | [2048, 64]  | rank:1/cuda:1\n",
            "small_table_0 | [0, 0]        | [512, 64]   | rank:0/cuda:0\n",
            "small_table_0 | [512, 0]      | [512, 64]   | rank:1/cuda:1\n",
            "small_table_1 | [0, 0]        | [512, 64]   | rank:0/cuda:0\n",
            "small_table_1 | [512, 0]      | [512, 64]   | rank:1/cuda:1\n",
            "rank:1,sharding plan: module: \n",
            "\n",
            "    param     | sharding type | compute kernel | ranks \n",
            "------------- | ------------- | -------------- | ------\n",
            "large_table_0 | row_wise      | fused          | [0, 1]\n",
            "large_table_1 | row_wise      | fused          | [0, 1]\n",
            "small_table_0 | row_wise      | fused          | [0, 1]\n",
            "small_table_1 | row_wise      | fused          | [0, 1]\n",
            "\n",
            "    param     | shard offsets | shard sizes |   placement  \n",
            "------------- | ------------- | ----------- | -------------\n",
            "large_table_0 | [0, 0]        | [2048, 64]  | rank:0/cuda:0\n",
            "large_table_0 | [2048, 0]     | [2048, 64]  | rank:1/cuda:1\n",
            "large_table_1 | [0, 0]        | [2048, 64]  | rank:0/cuda:0\n",
            "large_table_1 | [2048, 0]     | [2048, 64]  | rank:1/cuda:1\n",
            "small_table_0 | [0, 0]        | [512, 64]   | rank:0/cuda:0\n",
            "small_table_0 | [512, 0]      | [512, 64]   | rank:1/cuda:1\n",
            "small_table_1 | [0, 0]        | [512, 64]   | rank:0/cuda:0\n",
            "small_table_1 | [512, 0]      | [512, 64]   | rank:1/cuda:1\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "[rank0]:[W523 02:45:29.116195809 ProcessGroupNCCL.cpp:1250] Warning: WARNING: process group has NOT been destroyed before we destruct ProcessGroupNCCL. On normal program exit, the application should call destroy_process_group to ensure that any pending NCCL operations have finished in this process. In rare cases this process can exit before this point and block the progress of another member of the process group. This constraint has always been present,  but this warning has only been added since PyTorch 2.4 (function operator())\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "exit code: 0\n",
            "exit code: 0\n"
          ]
        }
      ],
      "source": [
        "spmd_sharing_simulation(ShardingType.ROW_WISE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TyJeUphwP3ut"
      },
      "source": [
        "Column-wise on the other hand, address the load imbalance problems for tables with large embedding dimensions. We will split the table vertically. Users can see that in the `shard_sizes` section in the printed plan log, the tables are halved by embedding dimension to be distributed onto two GPUs.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "vJ1lurhrP3ut"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "start for rank: 0\n",
            "start for rank: 1\n",
            "rank:0,sharding plan: module: \n",
            "\n",
            "    param     | sharding type | compute kernel | ranks\n",
            "------------- | ------------- | -------------- | -----\n",
            "large_table_0 | column_wise   | fused          | [0]  \n",
            "large_table_1 | column_wise   | fused          | [1]  \n",
            "small_table_0 | column_wise   | fused          | [0]  \n",
            "small_table_1 | column_wise   | fused          | [1]  \n",
            "\n",
            "    param     | shard offsets | shard sizes |   placement  \n",
            "------------- | ------------- | ----------- | -------------\n",
            "large_table_0 | [0, 0]        | [4096, 64]  | rank:0/cuda:0\n",
            "large_table_1 | [0, 0]        | [4096, 64]  | rank:1/cuda:1\n",
            "small_table_0 | [0, 0]        | [1024, 64]  | rank:0/cuda:0\n",
            "small_table_1 | [0, 0]        | [1024, 64]  | rank:1/cuda:1\n",
            "rank:1,sharding plan: module: \n",
            "\n",
            "    param     | sharding type | compute kernel | ranks\n",
            "------------- | ------------- | -------------- | -----\n",
            "large_table_0 | column_wise   | fused          | [0]  \n",
            "large_table_1 | column_wise   | fused          | [1]  \n",
            "small_table_0 | column_wise   | fused          | [0]  \n",
            "small_table_1 | column_wise   | fused          | [1]  \n",
            "\n",
            "    param     | shard offsets | shard sizes |   placement  \n",
            "------------- | ------------- | ----------- | -------------\n",
            "large_table_0 | [0, 0]        | [4096, 64]  | rank:0/cuda:0\n",
            "large_table_1 | [0, 0]        | [4096, 64]  | rank:1/cuda:1\n",
            "small_table_0 | [0, 0]        | [1024, 64]  | rank:0/cuda:0\n",
            "small_table_1 | [0, 0]        | [1024, 64]  | rank:1/cuda:1\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "[rank0]:[W523 02:45:37.247471296 ProcessGroupNCCL.cpp:1250] Warning: WARNING: process group has NOT been destroyed before we destruct ProcessGroupNCCL. On normal program exit, the application should call destroy_process_group to ensure that any pending NCCL operations have finished in this process. In rare cases this process can exit before this point and block the progress of another member of the process group. This constraint has always been present,  but this warning has only been added since PyTorch 2.4 (function operator())\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "exit code: 0\n",
            "exit code: 0\n"
          ]
        }
      ],
      "source": [
        "spmd_sharing_simulation(ShardingType.COLUMN_WISE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0kapPlSFP3ut"
      },
      "source": [
        "For `table-row-wise`, unfortuately we cannot simulate it due to its nature of operating under multi-host setup. We will present a python [`SPMD`](https://en.wikipedia.org/wiki/SPMD) example in the future to train models with `table-row-wise`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7HOhFal4P3ut"
      },
      "source": [
        "\n",
        "With data parallel, we will repeat the tables for all devices.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "ePzMsOp-P3uu"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "start for rank: 0\n",
            "start for rank: 1\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Sharding Type is data_parallel, caching params will be ignored\n",
            "Sharding Type is data_parallel, caching params will be ignored\n",
            "Sharding Type is data_parallel, caching params will be ignored\n",
            "Sharding Type is data_parallel, caching params will be ignored\n",
            "Sharding Type is data_parallel, caching params will be ignored\n",
            "Sharding Type is data_parallel, caching params will be ignored\n",
            "Sharding Type is data_parallel, caching params will be ignored\n",
            "Sharding Type is data_parallel, caching params will be ignored\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "rank:0,sharding plan: module: \n",
            "\n",
            "    param     | sharding type | compute kernel | ranks \n",
            "------------- | ------------- | -------------- | ------\n",
            "large_table_0 | data_parallel | dense          | [0, 1]\n",
            "large_table_1 | data_parallel | dense          | [0, 1]\n",
            "small_table_0 | data_parallel | dense          | [0, 1]\n",
            "small_table_1 | data_parallel | dense          | [0, 1]\n",
            "\n",
            "\n",
            "\n",
            "rank:1,sharding plan: module: \n",
            "\n",
            "    param     | sharding type | compute kernel | ranks \n",
            "------------- | ------------- | -------------- | ------\n",
            "large_table_0 | data_parallel | dense          | [0, 1]\n",
            "large_table_1 | data_parallel | dense          | [0, 1]\n",
            "small_table_0 | data_parallel | dense          | [0, 1]\n",
            "small_table_1 | data_parallel | dense          | [0, 1]\n",
            "\n",
            "\n",
            "\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "[rank0]:[W523 02:45:45.339723036 ProcessGroupNCCL.cpp:1250] Warning: WARNING: process group has NOT been destroyed before we destruct ProcessGroupNCCL. On normal program exit, the application should call destroy_process_group to ensure that any pending NCCL operations have finished in this process. In rare cases this process can exit before this point and block the progress of another member of the process group. This constraint has always been present,  but this warning has only been added since PyTorch 2.4 (function operator())\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "exit code: 0\n",
            "exit code: 0\n"
          ]
        }
      ],
      "source": [
        "spmd_sharing_simulation(ShardingType.DATA_PARALLEL)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
